#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简化系统推理脚本 - 使用系统Python环境
适用于没有完整PyTorch环境的情况
支持基础目标检测展示
"""

import cv2
import numpy as np
import time
import argparse
import os

def create_demo_detector():
    """创建演示检测器 - 使用简单的运动检测作为替代"""
    
    class DemoDetector:
        def __init__(self):
            self.background_subtractor = cv2.createBackgroundSubtractorMOG2(
                detectShadows=True
            )
            self.kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
            
        def detect(self, frame):
            """检测运动目标"""
            # 背景减法
            fg_mask = self.background_subtractor.apply(frame)
            
            # 形态学操作
            fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_OPEN, self.kernel)
            fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_CLOSE, self.kernel)
            
            # 查找轮廓
            contours, _ = cv2.findContours(
                fg_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
            )
            
            detections = []
            for contour in contours:
                area = cv2.contourArea(contour)
                if area > 500:  # 过滤小目标
                    x, y, w, h = cv2.boundingRect(contour)
                    confidence = min(area / 10000, 1.0)  # 简单的置信度计算
                    
                    detections.append({
                        'bbox': [x, y, x + w, y + h],
                        'confidence': confidence,
                        'class_name': 'Moving Object'
                    })
            
            return detections
    
    return DemoDetector()

def draw_detections(frame, detections):
    """绘制检测结果"""
    for det in detections:
        x1, y1, x2, y2 = det['bbox']
        confidence = det['confidence']
        class_name = det['class_name']
        
        # 绘制边界框
        color = (0, 255, 0)
        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
        
        # 绘制标签
        label = f"{class_name}: {confidence:.2f}"
        
        # 计算文本大小
        (text_width, text_height), baseline = cv2.getTextSize(
            label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1
        )
        
        # 绘制背景矩形
        cv2.rectangle(
            frame, 
            (x1, y1 - text_height - 10), 
            (x1 + text_width, y1), 
            color, 
            -1
        )
        
        # 绘制文本
        cv2.putText(
            frame, 
            label, 
            (x1, y1 - 5), 
            cv2.FONT_HERSHEY_SIMPLEX, 
            0.6, 
            (0, 0, 0), 
            1
        )
    
    return frame

def main():
    parser = argparse.ArgumentParser(description='简化系统推理演示')
    parser.add_argument('--camera', type=int, default=0, help='摄像头ID')
    parser.add_argument('--demo-mode', action='store_true', 
                       help='使用演示模式 (运动检测)')
    
    args = parser.parse_args()
    
    print("🚀 简化推理系统")
    print("=" * 40)
    print(f"📦 OpenCV版本: {cv2.__version__}")
    
    # 检查CUDA支持
    try:
        cuda_devices = cv2.cuda.getCudaEnabledDeviceCount()
        print(f"🚀 CUDA设备数: {cuda_devices}")
    except:
        print("💻 使用CPU模式")
    
    print(f"📷 摄像头ID: {args.camera}")
    
    # 检查是否有TensorRT引擎文件
    engine_path = "weights/yolo11n.engine"
    if os.path.exists(engine_path):
        print(f"⚠️  发现TensorRT引擎: {engine_path}")
        print("❌ 当前环境无法直接使用TensorRT引擎")
        print("💡 切换到演示模式 (运动检测)")
        args.demo_mode = True
    else:
        print("📁 未发现AI模型文件")
        args.demo_mode = True
    
    if args.demo_mode:
        print("🎯 使用演示模式: 运动目标检测")
        detector = create_demo_detector()
    else:
        print("❌ 无可用的推理模型")
        return
    
    # 打开摄像头
    print(f"\n📷 打开摄像头 {args.camera}...")
    cap = cv2.VideoCapture(args.camera)
    
    if not cap.isOpened():
        print(f"❌ 无法打开摄像头 {args.camera}")
        return
    
    # 设置摄像头参数
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    cap.set(cv2.CAP_PROP_FPS, 30)
    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
    
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    
    print(f"📸 摄像头配置: {width}x{height} @ {fps}FPS")
    print("\n🎮 控制说明:")
    print("  - 按 'q' 键退出")
    print("  - 按 's' 键截图")
    print("  - 按 'f' 键切换FPS显示")
    print("  - 按 'r' 键重置检测器")
    
    if args.demo_mode:
        print("  - 在摄像头前移动物体来测试检测")
    
    print("🎉 开始推理!\n")
    
    # 性能统计
    frame_count = 0
    start_time = time.time()
    show_fps = True
    inference_times = []
    
    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                print("❌ 无法读取摄像头帧")
                break
            
            frame_count += 1
            
            # 推理
            inference_start = time.time()
            detections = detector.detect(frame)
            inference_time = (time.time() - inference_start) * 1000
            inference_times.append(inference_time)
            
            # 绘制结果
            frame = draw_detections(frame, detections)
            
            # 显示信息
            info_text = f"检测: {len(detections)}个目标"
            cv2.putText(frame, info_text, (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            
            # 显示模式
            mode_text = "模式: 运动检测演示" if args.demo_mode else "模式: AI推理"
            cv2.putText(frame, mode_text, (10, 60), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1)
            
            # FPS信息
            if show_fps and frame_count % 30 == 0:
                current_time = time.time()
                elapsed = current_time - start_time
                fps = frame_count / elapsed if elapsed > 0 else 0
                avg_inference = np.mean(inference_times[-30:]) if inference_times else 0
                
                fps_text = f"FPS: {fps:.1f} | 处理: {avg_inference:.1f}ms"
                cv2.putText(frame, fps_text, (10, height - 20), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
            
            # 显示帧
            cv2.imshow('简化推理系统', frame)
            
            # 处理按键
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                print("👋 用户退出")
                break
            elif key == ord('s'):
                screenshot_name = f"simple_inference_{int(time.time())}.jpg"
                cv2.imwrite(screenshot_name, frame)
                print(f"📸 截图保存: {screenshot_name}")
            elif key == ord('f'):
                show_fps = not show_fps
                print(f"🔄 FPS显示: {'开启' if show_fps else '关闭'}")
            elif key == ord('r') and args.demo_mode:
                detector = create_demo_detector()
                print("🔄 检测器已重置")
    
    except KeyboardInterrupt:
        print("\n⏹️  收到中断信号")
    
    except Exception as e:
        print(f"\n❌ 运行错误: {e}")
    
    finally:
        # 清理资源和统计
        total_time = time.time() - start_time
        avg_fps = frame_count / total_time if total_time > 0 else 0
        avg_inference = np.mean(inference_times) if inference_times else 0
        
        print(f"\n📊 运行统计:")
        print(f"总运行时间: {total_time:.1f}s")
        print(f"处理帧数: {frame_count}")
        print(f"平均FPS: {avg_fps:.2f}")
        print(f"平均处理时间: {avg_inference:.2f}ms")
        
        cap.release()
        cv2.destroyAllWindows()
        print("🎉 程序结束")

if __name__ == "__main__":
    main()